#!/usr/bin/env python3

import pathlib
import random
from shapely.geometry import Polygon, GeometryCollection, Point
import jupedsim as jps

# Model parameters
NUM_FLOORS = 6
PEOPLE_PER_FLOOR = 20
FLOOR_SIZE = (20.0, 20.0)  # width, depth of each floor
STAIR_WIDTH = 2.0         # width of each staircase
FLOOR_GAP = 0.0           # no gap to keep geometry connected
AGENT_RADIUS = 0.3        # radius for each agent

# Movement parameters
DEFAULT_FREE_SPEED = 1.2  # m/s initial free speed
MIN_SPAWN_DIST = 2 * AGENT_RADIUS  # minimum distance between agents at spawn

# Logging callbacks (optional)
jps.set_debug_callback(lambda msg: print(f"DEBUG: {msg}"))
jps.set_info_callback(lambda msg: print(f"INFO: {msg}"))
jps.set_warning_callback(lambda msg: print(f"WARNING: {msg}"))
jps.set_error_callback(lambda msg: print(f"ERROR: {msg}"))

def create_building_geometry():
    """
    Build a 2D representation: place floors side-by-side on the x-axis,
    and add stair polygons connecting floors. Return a single connected
    GeometryCollection plus lists of floor and stair polygons.
    """
    floor_polys = []
    stair_polys = []
    for level in range(NUM_FLOORS):
        x0 = level * (FLOOR_SIZE[0] + FLOOR_GAP)
        # Floor rectangle
        floor_polys.append(Polygon([
            (x0, 0.0),
            (x0 + FLOOR_SIZE[0], 0.0),
            (x0 + FLOOR_SIZE[0], FLOOR_SIZE[1]),
            (x0, FLOOR_SIZE[1])
        ]))
        # Staircase down to previous floor
        if level > 0:
            for offset in [FLOOR_SIZE[0] / 3, 2 * FLOOR_SIZE[0] / 3]:
                sx = x0 + offset
                stair_polys.append(Polygon([
                    (sx, 0.0),
                    (sx + STAIR_WIDTH, 0.0),
                    (sx + STAIR_WIDTH, FLOOR_SIZE[1]),
                    (sx, FLOOR_SIZE[1])
                ]))
    # Merge all into one geometry
    return GeometryCollection(floor_polys + stair_polys), floor_polys, stair_polys


def sample_point_in_polygon(poly: Polygon) -> Point:
    """
    Uniformly sample a point inside the given polygon via rejection sampling.
    """
    minx, miny, maxx, maxy = poly.bounds
    while True:
        pt = Point(random.uniform(minx, maxx), random.uniform(miny, maxy))
        if poly.contains(pt):
            return pt


def main():
    # Build geometry
    geometry, floor_polys, stair_polys = create_building_geometry()

    # Initialize JuPedSim simulation
    sim = jps.Simulation(
        model=jps.CollisionFreeSpeedModel(),
        geometry=geometry,
        trajectory_writer=jps.SqliteTrajectoryWriter(output_file=pathlib.Path("evacuation_out.sqlite"))
    )

    # Add waypoint stages at centroids of stair polygons
    stair_stages = []
    for stair in stair_polys:
        cx, cy = stair.centroid.x, stair.centroid.y
        stair_stages.append(sim.add_waypoint_stage((cx, cy), STAIR_WIDTH / 2))

    # Add exit stages on ground (first) floor: small rectangles along top boundary
    exit_stages = []
    ground = floor_polys[0]
    minx, miny, maxx, maxy = ground.bounds
    exit_height = AGENT_RADIUS  # inset equal to agent radius
    for i in range(3):
        ex = minx + (i + 1) * (FLOOR_SIZE[0] / 4)
        exit_poly = [
            (ex,      maxy - exit_height),
            (ex + 2.0, maxy - exit_height),
            (ex + 2.0, maxy),
            (ex,      maxy)
        ]
        exit_stages.append(sim.add_exit_stage(exit_poly))

    # Define journey: from top stair descending all stair stages, then to an exit
    journey = jps.JourneyDescription()
    for stage in reversed(stair_stages):
        journey.add(stage)
    journey.add(exit_stages[0])
    journey_id = sim.add_journey(journey)

    # Agent parameter template
    agent_params = jps.CollisionFreeSpeedModelAgentParameters(
        journey_id=journey_id,
        stage_id=stair_stages[-1],
        radius=AGENT_RADIUS
    )
    agent_params.free_speed = DEFAULT_FREE_SPEED

    # Compute safe floor polygons by shrinking inward by agent radius
    safe_floor_polys = []
    for poly in floor_polys:
        safe = poly.buffer(-AGENT_RADIUS)
        if safe.is_empty:
            safe = poly  # fallback if buffer too large
        safe_floor_polys.append(safe)

    # Spawn agents uniformly on each safe floor polygon, avoiding close spawns
    all_spawned = []  # list of Point
    for safe_poly in safe_floor_polys:
        floor_spawned = []
        for _ in range(PEOPLE_PER_FLOOR):
            # find a valid spawn point
            while True:
                pt = sample_point_in_polygon(safe_poly)
                if all(pt.distance(existing) >= MIN_SPAWN_DIST for existing in floor_spawned):
                    floor_spawned.append(pt)
                    break
            agent_params.position = (pt.x, pt.y)
            sim.add_agent(agent_params)
            all_spawned.append(pt)

    # Run simulation until all evacuated
    print("Running evacuation simulation...")
    while sim.agent_count() > 0:
        sim.iterate()
    print(f"Evacuation complete in {sim.iteration_count()} iterations.")

if __name__ == "__main__":
    main()
